Amazon Titan Text Embedding V2 を試してみた

つくぼし

2024.05.06

こんにちは、つくぼし(tsukuboshi0755)です！

Amazon Titan Text EmbeddingsのV2が、Amazon Bedrockの新規埋め込みモデルとして利用可能になりました。

今回はこちらのAmazon Titan Text Embeddings V2を試してみたので紹介します！

G1とV2の比較

G1とV2の比較表は以下の通りです。

特徴	Titan Text Embeddings G1	Titan Text Embeddings V2
最大トークン数	8k	8k
ベクトルサイズ	1536	256/512/1024
対応言語数	25以上	100以上
正規化のサポート	なし	あり
100万トークンあたりの価格	0.10USD	0.02USD

V2はG1に比べて以下のようなメリットがあります。

より小さいベクトルサイズを選択する事で、データベースから文書抽出を検索して取得するための待ち時間を短縮可能
ベクトルの類似性を測定する際の精度の向上に役立つ、ベクトルの正規化をサポート
より幅広い言語がサポートされ、多言語の文書を埋め込む際に役立つ
トークンあたりの価格が低下し、さらにコストを削減可能

なお2024/5/6現在、Amazon Titan Text Embeddings V2はナレッジベースの埋め込みモデルとしては使用できないためご注意ください。

試してみた

今回はPythonのboto3を用いてTitan Text V2のAPIを呼び出し、2つの文を入力した上で、2つの文章の類似性を評価するコサイン類似度を計算してみます。

コサイン類似度については、以下のブログが詳しいので別途ご参照ください。

パラメータ

V2ではG1と比較して、APIで指定可能なリクエストボディのパラメータが増えているため、以下で紹介します。

パラメータ	指定可能な値	デフォルト値	用途
inputText	(文字列)	-	埋め込みベクトルに変換する入力テキスト。
normalize	true/false	true	埋め込みベクトルを正規化するかどうかを示すフラグ。 RAGで使用する場合はtrueでの設定が推奨される。
dimensions	256/512/1024	1024	埋め込みベクトルが持つ次元の数。待ち時間を短縮したい場合は小さく、逆に精度を高めたい場合は大きい値を指定すると良い。

特にV2ではベクトルサイズ(dimensions)を柔軟に選択できるようになったため、用途に応じて適切な値を選択すると良さそうです。

使用ライブラリ

今回使用するライブラリは、boto3のみです。

boto3：1.34.61

コサイン類似度を計算する際、一般的には以下のブログのようにNumPyやSciPyという外部ライブラリを別途導入し利用する事が多いです。

今回はあえて外部ライブラリを使わず、boto3以外はPythonの標準ライブラリのみでコサイン類似度を計算します。

boto3がプリインストールされているCloudShell等の環境であれば追加の外部ライブラリなしで実行できるので、ぜひそちらでお試しください。

スクリプト内容

Pythonのスクリプトを以下に示します。

スクリプト実行時の引数に比較対象となる2つの文を指定して実行すると、コサイン類似度が出力されます。

また今回はオレゴンジージョン(us-west-2)のBedrock呼び出しを想定しています。

import boto3
import json
import math
import sys

# メイン処理
def main():
    # 引数から比較対象となる2つの文を取得
    sentence1 = sys.argv[1]
    sentence2 = sys.argv[2]

    # 2つの文の埋め込みベクトルを計算
    vector1 = calculate_vector(sentence1)
    vector2 = calculate_vector(sentence2)

    # 2つの埋め込みベクトルのコサイン類似度を計算して出力
    similarity = cosine_similarity(vector1, vector2)
    print(similarity)

# テキストを入力として受け取り、埋め込みモデルを呼び出してベクトルを計算する関数
def calculate_vector(sentence):
    # Bedrock Runtimeクライアントを作成
    bedrock_runtime = boto3.client(
        service_name="bedrock-runtime", region_name="us-west-2"
    )

    # ベクトルサイズを256/512/1024のいずれかで指定
    dimensions = 256

    # リクエストボディを作成
    body = json.dumps(
        {
            "inputText": sentence,
            "dimensions": dimensions,
        }
    )

    # Bedrockの埋め込みモデルIDを指定
    modelId = "amazon.titan-embed-text-v2:0"

    # モデルを呼び出してレスポンスを取得
    response = bedrock_runtime.invoke_model(
        body=body,
        modelId=modelId,
        accept="*/*",
        contentType="application/json",
    )

    # レスポンスボディを取得
    response_body = json.loads(response.get("body").read())

    # 埋め込みベクトルを取得して返す
    embedding = response_body.get("embedding")
    return embedding

# ベクトルの内積を計算する関数
def dot_product(vecA, vecB):
    return sum(x * y for x, y in zip(vecA, vecB))

# ベクトルのユークリッドノルム(大きさ）を計算する関数
def magnitude(vec):
    return math.hypot(*vec)

# コサイン類似度を計算する関数
def cosine_similarity(vec1, vec2):
    return dot_product(vec1, vec2) / (magnitude(vec1) * magnitude(vec2))

if __name__ == "__main__":
    main()

実行結果

スクリプト実行前に、対象リージョンにおけるBedrockコンソールのモデルアクセスより、Amazon Titan Text Embeddings V2を有効化する必要があるためご注意ください。

有効化後、CloudShell上でmain.pyを実行した結果を以下に示します。

$ python main.py "東京の醤油ラーメン" "東京の味噌ラーメン"
0.8402169065862879

$ python main.py "東京の醤油ラーメン" "札幌の味噌ラーメン"
0.49552747137073716

$ python main.py "東京の味噌ラーメン" "札幌の味噌ラーメン"
0.5408871354010766

作成したスクリプトを実行すると、上記の通り問題なくコサイン類似度が計算できているようです。

例えば「東京の醤油ラーメン」、「東京の味噌ラーメン」、「札幌の味噌ラーメン」の中では、「東京の醤油ラーメン」と「東京の味噌ラーメン」のコサイン類似度が最も高くなっていることが確認できます。